package main

import (
	"context"
	"fmt"
	"net"
	"time"

	"atomgit.com/openlinksaas/proto-gen-go.git/net_proxy_api"
)

func createConnId(tunnelId string, client net_proxy_api.NetProxyApiClient) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	res, err := client.CreateConn(ctx, &net_proxy_api.CreateConnRequest{
		TunnelId: tunnelId,
	})
	if err != nil {
		return "", err
	}
	if res.Code != net_proxy_api.CreateConnResponse_CODE_OK {
		return "", fmt.Errorf(res.ErrMsg)
	}
	return res.ConnId, nil
}

func listenTcp(tunnelId, addrStr string, client net_proxy_api.NetProxyApiClient) error {
	addr, err := net.ResolveTCPAddr("tcp", addrStr)
	if err != nil {
		return err
	}
	li, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return err
	}
	go func() {
		for {
			time.Sleep(10 * time.Second)
			keepAlive(tunnelId, client)
		}
	}()
	for {
		conn, err := li.Accept()
		if err != nil {
			return err
		}
		connId, err := createConnId(tunnelId, client)
		if err != nil {
			fmt.Println(err)
			conn.Close()
			continue
		}
		go readTcpConn(tunnelId, connId, conn, client)
		go writeTcpConn(tunnelId, connId, conn, client)
	}
}

func readTcpConn(tunnelId, connId string, conn net.Conn, client net_proxy_api.NetProxyApiClient) {
	defer func() {
		recover()
	}()

	defer conn.Close()

	buf := make([]byte, 64*1024)
	for {
		n, err := conn.Read(buf)
		if n > 0 {
			err2 := sendData(tunnelId, connId, buf[:n], client)
			if err2 != nil {
				fmt.Println(err2)
				return
			}
		}
		if err != nil {
			fmt.Println(err)
			return
		}
	}
}

func writeTcpConn(tunnelId, connId string, conn net.Conn, client net_proxy_api.NetProxyApiClient) {
	defer func() {
		recover()
	}()

	defer conn.Close()

	for {
		data, err := recvData(tunnelId, connId, client)
		if err != nil {
			fmt.Println(err)
			return
		}
		writed := 0
		for writed < len(data) {
			n, err := conn.Write(data[writed:])
			if err != nil {
				fmt.Println(err)
				return
			}
			writed += n
		}
	}
}
